package com.example.deepai.service;

import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.DataType;
import io.milvus.param.collection.CreateCollectionParam;
import io.milvus.param.collection.FieldType;
import io.milvus.param.collection.HasCollectionParam;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.grpc.SearchResults;
import io.milvus.grpc.SearchResultData;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import com.example.deepai.model.UIElement;
import java.util.*;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

// 添加必要的导入
import jakarta.annotation.PostConstruct;
import io.milvus.param.collection.ShowCollectionsParam;
import io.milvus.grpc.ShowCollectionsResponse;  // 添加这个导入
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.DependsOn;

@Slf4j
@Service
@DependsOn("milvusConfig")
public class VectorDBService {
    private final MilvusServiceClient milvusClient;
    private static final String COLLECTION_NAME = "ui_components";
    private static final String ANALYSIS_COLLECTION_NAME = "ui_analysis";
    
    @Value("${milvus.host:localhost}")
    private String host;
    
    @Value("${milvus.port:19530}")
    private int port;

    public VectorDBService(MilvusServiceClient milvusClient) {
        this.milvusClient = milvusClient;
        // 移除构造函数中的初始化调用，改为使用 @PostConstruct
    }
    
    @PostConstruct
    public void init() {
        try {
            // 确保 Milvus 服务连接正常
            waitForMilvusReady();
            
            // 初始化默认集合
            initializeCollection();
            
            // 初始化分析集合
            initializeAnalysisCollection();
            
        } catch (Exception e) {
            log.error("Failed to initialize VectorDBService: {}", e.getMessage(), e);
        }
    }
    
    private void waitForMilvusReady() {
        int maxRetries = 5;
        int retryInterval = 5000; // 5秒
        
        for (int i = 0; i < maxRetries; i++) {
            try {
                HasCollectionParam hasParam = HasCollectionParam.newBuilder()
                    .withCollectionName(COLLECTION_NAME)
                    .build();
                milvusClient.hasCollection(hasParam);
                log.info("Successfully connected to Milvus server");
                return;
            } catch (Exception e) {
                log.warn("Waiting for Milvus server to be ready... Attempt {}/{}", i + 1, maxRetries);
                try {
                    Thread.sleep(retryInterval);
                } catch (InterruptedException ie) {
                    Thread.currentThread().interrupt();
                    break;
                }
            }
        }
        log.error("Failed to connect to Milvus server after {} attempts", maxRetries);
    }

    private void initializeCollection() {
        try {
            HasCollectionParam hasParam = HasCollectionParam.newBuilder()
                .withCollectionName(COLLECTION_NAME)
                .build();
            
            R<Boolean> response = milvusClient.hasCollection(hasParam);
            if (response != null && response.getData() != null && response.getData()) {
                log.info("Collection {} already exists", COLLECTION_NAME);
                return;
            }

            FieldType idField = FieldType.newBuilder()
                .withName("id")
                .withDataType(DataType.Int64)
                .withPrimaryKey(true)
                .withAutoID(true)
                .build();
                
            FieldType typeField = FieldType.newBuilder()
                .withName("type")
                .withDataType(DataType.VarChar)
                .withMaxLength(100)
                .build();
                
            FieldType descField = FieldType.newBuilder()
                .withName("description")
                .withDataType(DataType.VarChar)
                .withMaxLength(500)
                .build();
                
            FieldType vectorField = FieldType.newBuilder()
                .withName("feature_vector")
                .withDataType(DataType.FloatVector)
                .withDimension(512)
                .build();
                
            CreateCollectionParam createParam = CreateCollectionParam.newBuilder()
                .withCollectionName(COLLECTION_NAME)
                .withDescription("UI组件特征向量集合")
                .withShardsNum(2)
                .addFieldType(idField)
                .addFieldType(typeField)
                .addFieldType(descField)
                .addFieldType(vectorField)
                .build();
            
            milvusClient.createCollection(createParam);
            log.info("Collection {} created successfully", COLLECTION_NAME);
            
        } catch (Exception e) {
            log.error("Failed to initialize collection: {}", e.getMessage(), e);
            // 不抛出异常，让服务继续运行
        }
    }
    
    private void initializeAnalysisCollection() {
        try {
            if (!collectionExists(ANALYSIS_COLLECTION_NAME)) {
                FieldType idField = FieldType.newBuilder()
                    .withName("id")
                    .withDataType(DataType.Int64)
                    .withPrimaryKey(true)
                    .withAutoID(true)
                    .build();

                FieldType typeField = FieldType.newBuilder()
                    .withName("type")
                    .withDataType(DataType.VarChar)
                    .withMaxLength(100)
                    .build();

                FieldType descField = FieldType.newBuilder()
                    .withName("description")
                    .withDataType(DataType.VarChar)
                    .withMaxLength(500)
                    .build();

                FieldType vectorField = FieldType.newBuilder()
                    .withName("embedding")
                    .withDataType(DataType.FloatVector)
                    .withDimension(512)
                    .build();

                CreateCollectionParam createParam = CreateCollectionParam.newBuilder()
                    .withCollectionName(ANALYSIS_COLLECTION_NAME)
                    .withDescription("UI components analysis collection")
                    .addFieldType(idField)
                    .addFieldType(typeField)
                    .addFieldType(descField)
                    .addFieldType(vectorField)
                    .withShardsNum(2)
                    .build();

                milvusClient.createCollection(createParam);
                log.info("Created {} collection", ANALYSIS_COLLECTION_NAME);
            }
        } catch (Exception e) {
            log.error("Failed to initialize analysis collection", e);
        }
    }

    private boolean collectionExists(String collectionName) {
        try {
            ShowCollectionsParam showParam = ShowCollectionsParam.newBuilder()
                .withCollectionNames(Collections.singletonList(collectionName))
                .build();
            // 修复: 正确获取集合名称列表
            R<ShowCollectionsResponse> response = milvusClient.showCollections(showParam);
            return response.getData().getCollectionNamesList().contains(collectionName);
        } catch (Exception e) {
            log.error("Failed to check if collection exists: {}", e.getMessage());
            return false;
        }
    }

    public void storeUIComponent(UIElement element, float[] vector) {
        try {
            // 转换为 List<Float>
            List<Float> vectorList = new ArrayList<>(vector.length);
            for (float v : vector) {
                vectorList.add(v);
            }

            List<InsertParam.Field> fields = new ArrayList<>();
            
            // 添加类型字段
            if (element.getType() != null) {
                // 修复: 使用正确的 Field 构造方法
                fields.add(new InsertParam.Field("type", Collections.singletonList(element.getType())));
            }
            
            // 添加描述字段 - 检查 UIElement 是否有 description 属性
            if (element.getText() != null) {  // 修改为使用 getText() 方法
                fields.add(new InsertParam.Field("description", Collections.singletonList(element.getText())));
            }
            
            // 添加向量字段
            fields.add(new InsertParam.Field("feature_vector", Collections.singletonList(vectorList)));

            InsertParam insertParam = InsertParam.newBuilder()
                .withCollectionName(COLLECTION_NAME)
                .withFields(fields)
                .build();

            milvusClient.insert(insertParam);
            log.info("UI component stored successfully: {}", element.getType());
        } catch (Exception e) {
            log.error("Failed to store UI component: {}", e.getMessage(), e);
            throw new RuntimeException("Failed to store UI component", e);
        }
    }
    
    public List<UIElement> findSimilarComponents(float[] queryVector) {
        try {
            // 将 float[] 转换为 List<Float>
            List<Float> vectorList = new ArrayList<>(queryVector.length);
            for (float v : queryVector) {
                vectorList.add(v);
            }
            
            // 构建搜索参数，使用 List<Float> 类型
            SearchParam searchParam = SearchParam.newBuilder()
                .withCollectionName(COLLECTION_NAME)
                .withMetricType(MetricType.L2)
                .withOutFields(Collections.singletonList("id"))
                .withVectors(Collections.singletonList(vectorList))  // 使用 List<Float> 而不是 float[]
                .withTopK(5)
                .build();
            
            // 执行搜索
            R<SearchResults> searchResponse = milvusClient.search(searchParam);
            SearchResults results = searchResponse.getData();
            
            // 转换结果为 UIElement 列表
            List<UIElement> similarElements = new ArrayList<>();
            SearchResultData searchResultData = results.getResults();
            List<Long> ids = new ArrayList<>();
            searchResultData.getIds().getIntId().getDataList().forEach(id -> 
                ids.add(Long.valueOf(id))
            );
            
            for (Long id : ids) {
                UIElement element = new UIElement();
                element.setId(id);
                similarElements.add(element);
            }
            
            return similarElements;
            
        } catch (Exception e) {
            log.error("Failed to find similar components: {}", e.getMessage(), e);
            return new ArrayList<>();
        }
    }
}